fix(reshard): make distributed refit reliable and faster - #635

What this PR does

This PR completes the ModelExpress half of the NeMo-RL reshard-refit path. It combines seven focused changes:

Each change below removes a specific source of wasted work or incorrect weight placement:

• Collapse duplicate replica offers before planning. Data-parallel and expert-data-parallel ranks often advertise identical copies of the same model shard. Keeping every copy made the receiver read the same bytes repeatedly and broke large reads into hundreds of thousands of small ones. This keeps one verified copy and leaves genuinely different shards untouched.
• Build NIXL read descriptors once for a stable plan. Turning a transfer plan into low-level NIXL read requests is setup work. Rebuilding those requests on every refit added CPU overhead even when neither the model layout nor the selected source ranks changed.
• Install reconstructed tensors in one batch. Full-pull tensors used to be copied into live parameters through thousands of individual Python/CUDA operations. Batching those copies preserves the same destinations while removing repeated launch and interpreter overhead.
• Recognize vLLM's fused MoE expert views. vLLM exposes grouped experts through tensor views such as unsqueeze(0).unbind(). ModelExpress previously treated that valid layout as unsupported, so BF16 MoE refits could not complete. The capture logic now records the real expert storage without copying it.
• Release registered memory before shutting NIXL down. Destroying a NIXL agent while GPU memory was still registered could abort an otherwise successful worker during teardown. Shutdown now releases those registrations in the required order.
• Check publisher readiness without decoding the whole model. A quorum check only needs the published training step and whether the table is non-empty. Parsing every tensor and shard just to answer that question delayed each refit without improving correctness.
• Map QKV using global row ranges when KV heads are below TP. Computing num_kv_heads // tp_size gives zero for Nemotron Ultra's 2 KV heads at TP8. Megatron actually splits the fused QKV tensor by raw rows, so the fix maps each rank's real row interval into Q, K, and V instead of inventing a local KV-head count.

The NeMo-RL publisher and lifecycle integration is in NVIDIA-NeMo/RL #3632. The KV<TP fix needs both PRs: NeMo-RL publishes global, per-layer QKV geometry, and ModelExpress maps each rank's raw fused-row interval.

Data flow

flowchart LR
    L[Live Megatron layer config] --> N[Global per-layer Q/KV descriptor]
    N --> P[Raw TP fused-row shard]
    P --> I[Global interval intersection]
    I --> Q[q_proj shards]
    I --> K[k_proj shards]
    I --> V[v_proj shards]

    A[Trainer shard tables] --> D[Deduplicate identical replicas]
    Q --> D
    K --> D
    V --> D
    D --> C[Build and cache read descriptors]
    C --> R[Batched NIXL reads]
    R --> B[Batched install]

    S[Publisher step stamps] --> U[Cheap quorum check]
    B --> X[Deregister memory]
    X --> Y[Destroy NIXL agent]

Megatron does not give every TP rank one KV head when KV heads are below TP. It slices the globally interleaved fused QKV tensor by raw rows, so most ranks legitimately publish Q only. The new interval path intersects each raw source interval with the global Q/K/V bands and merges the sparse per-rank offers into complete HF tensors.

Measured effect of the original six changes

The performance evidence below predates the new GQA commit and must not be copied to that arm.

On real Qwen3-30B BF16 tensors with Megatron EP8 publishers and a vLLM TP2 receiver, each receiver rank changed from:

• 47.58 GiB to 37.82 GiB on the wire;
• 809,112 to 19,011 read segments;
• 2791.3 ms to 469.5 ms median wire time;
• 5102.6 ms to 480.7 ms median receiver refit time;
• 54.5% to more than 97% stage attribution.

The complete pre-GQA branch was also tested at 32 GPUs (16 trainer and 16 receiver GPUs) over NIXL/RDMA. Three independent dense cold starts per arm produced a 7.57x receiver speedup. MoE reached 100% coverage across 11 refits, exact verification found 0 of 435 parameters changed on all 16 receiver ranks for a same-checkpoint refit, and moving-model GRPO probability/JS-divergence/non-zero-gradient gates passed.

GQA/KV-heads-below-TP evidence

The new QKV interval branch is currently unit-qualified plus representative CUDA-tensor-qualified:

• Q=64, KV=2, head_dim=128, logical trainer TP8 reconstructs byte-exact Q/K/V with no gaps or overlaps;
• ranks without K/V rows publish Q only;
• sparse tables merge into complete sources and build a bounded/full-pull plan with zero fallback;
• divisible layouts remain byte-for-byte identical to the legacy local-head path;
• 24Q/6KV/TP4 and heterogeneous per-layer geometry are covered;
• missing/invalid global geometry fails closed;
• BF16 CUDA tensors passed one same-weight refit and one changed-weight refit with exact Q/K/V parameters and projection-output equality.

This is not yet a full Megatron-to-vLLM TP8 E2E claim, and it is not full Nemotron Ultra qualification. The PR remains draft until the companion NeMo-RL head and real TP8 generator gate complete.

Review guide

1. Global QKV interval mapping (new correctness fix)

Commit: f719207 — fix(reshard): map global QKV intervals when KV heads are below TP

Review:

• modelexpress_rl/train/engines/megatron/aliases.py
• tests/test_reshard_megatron_gqa.py

Please check the half-specified metadata rejection, raw source interval validation, per-group overlap math, omission of empty K/V tensors, and exact-once source-row coverage.

2. Replica deduplication (main performance change)

Commit: c55d80c

Review refit/reshard/rendezvous.py and tests/test_reshard_refit_replica_merge.py. Only offers with identical geometry and digest should collapse.

3. Fused MoE capture (main pre-GQA correctness change)

Commit: bdc75f8

Review geometry.py, receiver.py, types.py, and test_reshard_refit_moe_experts.py, especially unsqueeze(0).unbind() and keyword-only loader calls.

4. NIXL lifecycle

Commit: a3382e2

Review nixl_transfer.py and test_nixl_peer_lifecycle.py. Every registered region must be released before agent destruction, including partial setup and repeated shutdown.

5. Batched install and descriptor caching

Commits: 6d86828, dbaa534

Review cache invalidation when the plan or wire arm changes, and verify batched copies preserve destination order and placement.

6. Quorum parse skip

Commit: 55aff9f

Review rendezvous.py and its tests. The fetch remains serial intentionally; concurrent metadata reads measured slower.

Scope still pending

• real Megatron TP8 to vLLM TP8 E2E for Q=64/KV=2, with exact receiver parameters and generation agreement;
• full 520-GPU Nemotron Ultra qualification;
• FP8 installation (separate meta-tensor issue);
• restart and elastic-scale qualification;
• the future O(1) quorum protocol change.

Test plan

• PYTHONPATH=. python3 -m pytest tests/test_*reshard*.py tests/test_*nixl*.py -q — 266 passed
• CPU TP8 reconstruction, rendezvous merge, planner, full-pull and negative controls
• BF16 CUDA same-weight and moving-weight parameter/projection parity
• Existing real 32-GPU BF16 dense/MoE gates for the original six commits
• Companion NeMo-RL #3632 CI on the global per-layer descriptor contract
• Real Megatron TP8 -> vLLM TP8 Q=64/KV=2 E2E
• Full Nemotron Ultra qualification
• CI and CodeRabbit on head f719207

Summary by CodeRabbit

• New Features
  • Added options to enable batched refit installation and descriptor caching, improving repeated resharding performance.
  • Improved rendezvous discovery with faster metadata handling, entry counts, replica deduplication, and resilience to individual fetch failures.
  • Added support for additional fused-MoE and grouped-query attention resharding layouts.
• Bug Fixes
  • Improved cleanup of registered memory during shutdown.
  • Added clearer reporting of unsupported refit operations and their causes.
• Documentation
  • Documented the new batched-install configuration option.